{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Simulating a MAST-U shot using input data from EFIT++ (via pyUDA)\n",
    "\n",
    "This example notebook will show how to carry out static forward Grad-Shafranov equilibrium solves for MAST-U shots using FreeGSNKE. This example will follow the methodology for solving the GS equation in FreeGSNKE laid out in [Pentland et al. (2024)](https://arxiv.org/abs/2407.12432).\n",
    "\n",
    "We will:\n",
    " - load the MAST-U machine description (from pyUDA).\n",
    " - load the input parameters from EFIT++ **magnetics-only** reconstructions of a given shot (inputs required for each time slice being simulated). These include the coil currents and profile functions/parameters (plus some other quantities). We'll do some input checking to ensure no spurious parameters exist. \n",
    " - simulate the shot using FreeGSNKE's static forward solver.\n",
    " - compare poloidal flux and shape control targets between FreeGSNKE and the EFIT++ reconstructions.\n",
    "\n",
    "To re-familiarise yourself with the static forward solver in FreeGSNKE, revisit the \"example1\" notebooks. \n",
    "\n",
    "**Note**: This notebook can only be run with access to a UDA client (or run via the Freia cluster at UKAEA). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from copy import deepcopy\n",
    "\n",
    "# freegsnke modules (other may be needed later)\n",
    "from freegsnke import equilibrium_update, GSstaticsolver, mastu_tools # the last module containing various functions for pulling MAST-U data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# these paths are required to be able to access the MAST-U data via the UDA module\n",
    "os.environ[\"UDA_HOST\"] = \"uda2.hpc.l\"\n",
    "# os.environ[\"UDA_HOST\"] = \"data.mastu.ukaea.uk\"\n",
    "os.environ[\"UDA_PORT\"] = \"56565\"\n",
    "os.environ[\"UDA_META_PLUGINNAME\"] = \"MASTU_DB\"\n",
    "os.environ[\"UDA_METANEW_PLUGINNAME\"] = \"MASTU_DB\"\n",
    "\n",
    "# initialise the pyuda client to access the data\n",
    "import pyuda\n",
    "client = pyuda.Client()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select MAST-U shot that you wish to simulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MAST-U shot number (make sure it's a valid shot with EFIT++ equilibrium reconstruction data!)\n",
    "shot = 45425"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the MAST-U machine\n",
    "\n",
    "Next, we build the MAST-U machine object. This can be done using the geometry data from a particular shot number via the 'get_machine_data' function below.\n",
    "\n",
    "FreeGSNKE requires the following environment variables in order to build the machine: ACTIVE_COILS_PATH, PASSIVE_COILS_PATH, WALL_PATH, LIMITER_PATH, PROBE_PATH. \n",
    "\n",
    "Note 1: There are two options used for building the active coils in FreeGSNKE. The first sets up **symmetric** coils where the corresponding up/down coils in MAST-U (e.g. P4_upper and P4_lower) are wired **in series** and therefore will be assigned the same current value later. The second option sets up **non-symmetric** coils in which the up/down coils are modelled independently and can be assigned **different** currents (i.e. they are **not wired in series**). Set the 'symmetric_machine' parameter to either 'True' or 'False' to toggle this setting. We use the **non-symmetric** setting here as the coil currents in EFIT++ are calculated independently (though you may test out the symmetric case too to see the differences). \n",
    "\n",
    "Note 2: The magnetic probes are not used in this script (see \"example1c\" for how to use them). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function will pull machine geoemtry data from the shot number.\n",
    "# 'split_passive=True' means the passives structures will be modelled as parallelograms with current\n",
    "# distributed evenly over their cross-sectional area (which is reccomended), setting to 'False' will\n",
    "# model them as point sources (less accurate for parallelograms with large surface areas).\n",
    "\n",
    "# the first time this runs it may take a while (though if you have the machine description files already in the 'save_path' you can comment this out)\n",
    "mastu_tools.get_machine_data(save_path=\"../machine_configs/MAST-U\", shot=shot, split_passives=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set paths (choose symmetric or non-symmetric active coils)\n",
    "symmetric_machine = True\n",
    "if symmetric_machine:\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_active_coils.pickle\"\n",
    "else:\n",
    "    active_coils_path=\"../machine_configs/MAST-U/MAST-U_active_coils_nonsym.pickle\"\n",
    "passive_coils_path=\"../machine_configs/MAST-U/MAST-U_passive_coils.pickle\"\n",
    "limiter_path=\"../machine_configs/MAST-U/MAST-U_limiter.pickle\"\n",
    "wall_path=\"../machine_configs/MAST-U/MAST-U_wall.pickle\"\n",
    "magnetic_probe_path=\"../machine_configs/MAST-U/MAST-U_magnetic_probes.pickle\"\n",
    "\n",
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=active_coils_path,\n",
    "    passive_coils_path=passive_coils_path,\n",
    "    limiter_path=limiter_path,\n",
    "    wall_path=wall_path,\n",
    "    magnetic_probe_path=magnetic_probe_path,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the required input data from the EFIT++ reconstructions\n",
    "Here, we load the parameters required by the static forward GS solver (at each independent EFIT++ time step). They include the:\n",
    " - plasma current ($I_p$).\n",
    " - fvac parameter (= $r B_{tor}$).\n",
    " - profile function coefficients (alpha, beta - both vectors) and logical parameters (alpha_logic, beta_logic - both Booleans).\n",
    " - active coil and passive structure currents (loads currents for both symmetric and non-symmetric active coil setups).\n",
    "    \n",
    "Note 1: The function loads values for all time steps in the EFIT++ reconstruction (including time steps that may have failed to converge - we will exclude these later on).\n",
    "\n",
    "Note 2: For **magnetics-only** EFIT++ reconstructions, the \"Lao\" polynomials are used to parameterise the pressure and toroidal current density profiles, respectively:\n",
    "\n",
    "$$\n",
    "\\frac{dp}{d \\hat{\\psi}} = \\sum_{n=0}^{n_p} \\alpha_n \\hat{\\psi}^n - \\hat{\\alpha} \\hat{\\psi}^{n_p+1} \\sum_{n=0}^{n_p} \\alpha_n,\n",
    "$$\n",
    "\n",
    "$$\n",
    "F \\frac{dF}{d \\hat{\\psi}} = \\sum_{n=0}^{n_F} \\beta_n \\hat{\\psi}^n - \\hat{\\beta} \\hat{\\psi}^{n_F+1} \\sum_{n=0}^{n_F} \\beta_n.\n",
    "$$\n",
    "\n",
    "Where:\n",
    "-  $\\hat{\\psi}$ is the normalised flux.\n",
    "-  $ \\alpha = (\\alpha_0, \\ldots, \\alpha_{n_p})$ and $\\beta = (\\beta, \\ldots, \\beta{n_F})$ are the profile coefficients.\n",
    "-  $\\hat{\\alpha}$ and $\\hat{\\beta}$ are logical (Boolean) parameters that set the boundary condition for the functions.\n",
    "\n",
    "Note 3: There is also the option to set the passive structure currents to zero if desired. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the parameters\n",
    "Ip, fvac, alpha, beta, alpha_logic, beta_logic, currents_sym, currents_nonsym, _ = mastu_tools.load_static_solver_inputs(\n",
    "    client=client,        # pyuda client\n",
    "\tshot=shot,            # shot number\n",
    "\tzero_passives=False,  # set passive structure currents to zero? (True or False)\n",
    "    active_coils_path=active_coils_path,\n",
    "    passive_coils_path=passive_coils_path,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now load the shot times at which the EFIT++ reconstructions were carried out (which correspond to the above data) and check which time slices successfully converged. This is important because we need to exclude those that failed to converge (typically these occur during the ramp-up and ramp-down of the plasma). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the efit reconstruction timestamps and convergence status\n",
    "# Note: efit_times that have efit_status=-1 did not converge so we will exclude those time slices\n",
    "efit_times, efit_status = mastu_tools.load_efit_times_and_status(client=client, shot=shot)\n",
    "\n",
    "# figure out which time slices did not converge\n",
    "time_indices = np.where(efit_status == 1)[0]\n",
    "time_slices_excluded = np.where(efit_status == -1)[0]\n",
    "print(f\"{len(time_slices_excluded)} time slices (out of total {len(efit_times)}) excluded from simulation.\")\n",
    "print(f\"The excluded time slices are: {efit_times[time_slices_excluded]} seconds.\")\n",
    "\n",
    "times = efit_times[time_indices]  # these are the slices we simulate\n",
    "print(f\"Total time slices to be simulated: {len(times)}.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us plot some of the input parameters over these times to check they are indeed valid (as we may need to exclude some more if EFIT++ failed to catch the non-convergence)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ax1 = plt.subplots(1, 1, figsize=(12,4), dpi=80)\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, Ip[time_indices], color='k', linewidth=1, marker='x', markersize=2, zorder=10)\n",
    "ax1.set_xlabel(r'Shot time [$s$]')\n",
    "ax1.set_ylabel(r'$I_p$ [A]')\n",
    "ax1.set_xticks(np.arange(0, times[-1], 0.1))\n",
    "ax1.set_xlim(0, times[-1] + 0.005)\n",
    "ax1.ticklabel_format(axis='y', scilimits=(0,0))\n",
    "\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(12,4), dpi=80)\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, alpha[time_indices,0], color='b', linewidth=1, marker='x', markersize=2, zorder=10, label=r\"$\\alpha_0$\")\n",
    "ax1.plot(times, alpha[time_indices,1], color='r', linewidth=1, marker='x', markersize=2, zorder=10, label=r\"$\\alpha_1$\")\n",
    "ax1.set_xlabel(r'Shot time [$s$]')\n",
    "ax1.set_ylabel(r'$\\alpha$')\n",
    "ax1.set_xticks(np.arange(0, times[-1], 0.1))\n",
    "ax1.set_xlim(0, times[-1] + 0.005)\n",
    "ax1.ticklabel_format(axis='y', scilimits=(0,0))\n",
    "ax1.legend(loc = \"upper left\")\n",
    "\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(12,4), dpi=80)\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, beta[time_indices,0], color='b', linewidth=1, marker='x', markersize=2, zorder=10, label=r\"$\\beta_0$\")\n",
    "ax1.plot(times, beta[time_indices,1], color='r', linewidth=1, marker='x', markersize=2, zorder=10, label=r\"$\\beta_1$\")\n",
    "ax1.set_xlabel(r'Shot time [$s$]')\n",
    "ax1.set_ylabel(r'$\\beta$')\n",
    "ax1.set_xticks(np.arange(0, times[-1], 0.1))\n",
    "ax1.set_xlim(0, times[-1] + 0.005)\n",
    "ax1.ticklabel_format(axis='y', scilimits=(0,0))\n",
    "ax1.legend(loc = \"lower left\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the $\\beta$ profile parameters at the final time slice look quite spurious and so we can cut those off."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_indices = time_indices[0:-1]\n",
    "times = efit_times[time_indices]\n",
    "n = len(times)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Side example: plotting the profiles at a given time slice\n",
    "\n",
    "Here, we will plot the $p'$ and $FF'$ profiles at a given time slice to see what they look like. Note that here we're showing how they can be constructed from the alpha and beta parameters alone - once the `profiles` objects have been initialised later on, they can be plotted directly from them. \n",
    "\n",
    "First, we'll select a time slice and extract the parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = time_indices[110]\n",
    "times[index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_n = np.linspace(0,1,65)             # normalised flux from values 0 (at magnetic axis) to 1 (at plasma boundary)\n",
    "alpha_coeffs = alpha[index,:]           # select the alpha coefficients from the relevant time slice\n",
    "beta_coeffs = beta[index,:]             # do the same bor the beta coefficients\n",
    "alpha_logic_ = bool(alpha_logic[index])  # same for the logical parameters\n",
    "beta_logic_ =  bool(beta_logic[index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if the logical parameters are true, we need to add the final coefficients to alpha and beta (see the definitions above)\n",
    "if alpha_logic_:\n",
    "    alpha_coeffs = np.hstack([alpha_coeffs, -np.sum(alpha_coeffs)])\n",
    "if beta_logic_:\n",
    "    beta_coeffs = np.hstack([beta_coeffs, -np.sum(beta_coeffs)])\n",
    "print(f\"alpha = {alpha_coeffs}\")\n",
    "print(f\"beta = {beta_coeffs}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the exponents of psi\n",
    "exp_alpha = (np.tile(psi_n, (len(alpha_coeffs),1)).T)**range(0,len(alpha_coeffs))\n",
    "exp_beta = (np.tile(psi_n, (len(beta_coeffs),1)).T)**range(0,len(beta_coeffs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# finally, calculate the profiles for plotting below\n",
    "pprime = np.sum(alpha_coeffs*exp_alpha, axis=1)\n",
    "ffprime = np.sum(beta_coeffs*exp_beta, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(15,6), dpi=80)\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(psi_n, pprime, color='k', linewidth=1, marker='x', markersize=2, zorder=10, label=\"Reconstructed profile\")\n",
    "ax1.set_xlabel(r'$\\hat{\\psi}$')\n",
    "ax1.set_ylabel(r\"$p'(\\hat{\\psi})$\")\n",
    "ax1.set_xticks(np.arange(0, 1.05, 0.1))\n",
    "ax1.set_xlim(0, 1 + 0.005)\n",
    "ax1.ticklabel_format(axis='y', scilimits=(0,0))\n",
    "\n",
    "ax2.grid(zorder=0, alpha=0.75)\n",
    "ax2.plot(psi_n, ffprime, color='k', linewidth=1, marker='x', markersize=2, zorder=10, label=\"Reconstructed profile\")\n",
    "ax2.set_xlabel(r'$\\hat{\\psi}$')\n",
    "ax2.set_ylabel(r\"$FF'(\\hat{\\psi})$\")\n",
    "ax2.set_xticks(np.arange(0, 1.05, 0.1))\n",
    "ax2.set_xlim(0, 1 + 0.005)\n",
    "ax2.legend(loc=\"lower left\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see how different these profiles are to the ones in the next notebook where tension splines are used instead of polynomials."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialise the equilibrium object and static solver\n",
    "Here we initialise the equilibrium object as we have in the previous example scripts. We use the same rectangular computational grid and domain as used by EFIT++ for easy comparison later on (we could use a denser grid, however, we would need to interpolate some things later on). \n",
    "\n",
    "We also load FreeGSNKE's static GS solver (can be used for forward and inverse solves but here we use for forward only). \n",
    "\n",
    "Note 1: It is not necessary to instantiate a new solver object when using it to solve for multiple different equilibria, as long as the integration domain, mesh grid, and tokamak required are identical. If you wish to modify these values for the same/different equilibria then you will need to instantiate new objects. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# equilibrium object (note that both nx and ny have to be of the form 2**n + 1 with n being an integer)\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "tokamak=tokamak,             # sets up the object with the MAST-U tokamak\n",
    "Rmin=0.06, Rmax=2.0,         # computational grid radial limits (same as EFIT++)\n",
    "Zmin=-2.2, Zmax=2.2,         # computational grid vertical limits (same as EFIT)\n",
    "nx=65,                       # number of grid points in the radial direction\n",
    "ny=65,                       # number of grid points in the vertical direction\n",
    "psi=None                     # initial guess for the plasma flux (can provide one if available)\n",
    ")\n",
    "plasma_psi = deepcopy(eq.plasma_psi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# static solver object (used for solving later on)\n",
    "solver = GSstaticsolver.NKGSsolver(eq)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use FreeGSNKE to solve the static forward GS equation at each time slice\n",
    "\n",
    "We are finally ready to solve the static forward GS problem using FreeGSNKE. This will be done sequentially in a loop over the times that we wish to simulate (setting this up to run in parallel would be cool). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import jtor_update # for initialising the profile object\n",
    "from copy import deepcopy\n",
    "\n",
    "# carry out the simulation and store some quantities of interest\n",
    "freegsnke_eqs = []\n",
    "freegsnke_profiles = []\n",
    "freegsnke_runtimes = []\n",
    "for i in time_indices:\n",
    "\n",
    "    # strings for dictionary keys\n",
    "    time_str = str(efit_times[i]).replace(\".\", \"\")\n",
    "    name = f\"shot_{shot}_time_{time_str}\"\n",
    "\n",
    "    # start timer\n",
    "    start_time = time.time()\n",
    "\n",
    "    print(f\"-----Solving equilibrium for shot {shot} at time {efit_times[i]}-----\")\n",
    "    \n",
    "    # initialise profile object\n",
    "    profiles = jtor_update.Lao85(\n",
    "        eq=eq,                                          # equilibrium object\n",
    "        Ip=Ip[i],                                       # total plasma current\n",
    "        fvac=fvac[i],                                   # f vacuum parameter (R*Bt)\n",
    "        alpha=alpha[i,:],                               # p' profile coefficients\n",
    "        beta=beta[i,:],                                 # ff' profile coefficients\n",
    "        alpha_logic=bool(alpha_logic[i]),               # logic parameters from above\n",
    "        beta_logic=bool(beta_logic[i]),\n",
    "    )\n",
    "    \n",
    "    # set coil currents in eq object\n",
    "    # checks if machine is up/down symmetric or not\n",
    "    if symmetric_machine:\n",
    "        for key in currents_sym.keys():\n",
    "            eq.tokamak.set_coil_current(key, currents_sym[key][i])\n",
    "    else:\n",
    "        for key in currents_nonsym.keys():\n",
    "            eq.tokamak.set_coil_current(key, currents_nonsym[key][i])\n",
    "     \n",
    "    # carry out the forward solve\n",
    "    eq.set_plasma_psi(plasma_psi)\n",
    "    solver.solve(eq=eq,\n",
    "                 profiles=profiles,\n",
    "                 constrain=None,\n",
    "                 target_relative_tolerance=1e-6,\n",
    "                )\n",
    "\n",
    "    # stop timer\n",
    "    end_time = time.time()\n",
    "\n",
    "    # store data\n",
    "    freegsnke_eqs.append(deepcopy(eq))\n",
    "    freegsnke_profiles.append(deepcopy(profiles))\n",
    "    freegsnke_runtimes.append(end_time - start_time)\n",
    "\n",
    "\n",
    "print(\"-----SIMULATIONS COMPLETE!-----\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the FreeGSNKE runtimes might be slower than as reported in [Pentland et al. (2024)](https://arxiv.org/abs/2407.12432) because this script is run directly on Freia (which is slow!) rather than on a faster MacBook Pro machine..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Median FreeGSNKE runtime per time slice = {np.median(freegsnke_runtimes)} seconds.\")\n",
    "print(f\"Total FreeGSNKE runtime = {np.sum(freegsnke_runtimes)} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare FreeGSNKE simulations with EFIT++ reconstructions\n",
    "Here, we can compare some of the quantities produced by the FreeGSNKE simulations and the EFIT++ reconstructions. \n",
    "\n",
    "How to access some of the important quantities of interest in FreeGSNKE are listed below (there are more!):\n",
    "- \"psi_total\" = eq.psi()\n",
    "- \"psi_plasma\" = eq.plasma_psi\n",
    "- \"psi_coils\" = eq.tokamak_psi\n",
    "- \"psi_axis\" = eq.psi_axis\n",
    "- \"psi_boundary\" = eq.psi_bndry\n",
    "- \"Jtor\" = profiles.jtor\n",
    "- \"magnetic_axis\" = eq.magneticAxis()[0:2]\n",
    "- \"midplane_inner_outer_radii\" = eq.innerOuterSeparatrix(0)\n",
    "- \"x_points\" = eq.xpt[:,0:2]\n",
    "- \"pprime\" = eq.pprime(np.linspace(0,1,65))/profiles.L\n",
    "- \"ffprime\" = eq.ffprime(np.linspace(0,1,65))/profiles.L\n",
    "- \"strike_points\" = find_strikepoints(R=eq.R, Z=eq.Z, psi_total=eq.psi(), psi_boundary=eq.psi_bndry, limiter=np.array([eq.tokamak.limiter.R,eq.tokamak.limiter.Z]).T)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, however, we need to extract the EFIT++ output data that we want to compare. This is done via the 'extract_EFIT_outputs' function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "efit_psi_total, efit_psi_axis, efit_psi_boundary, efit_jtor, efit_magnetic_axis, efit_midplane_inner_outer_radii, efit_x_points, efit_pprime, efit_ffprime, efit_strike_points, efit_fluxloops, efit_pickups = mastu_tools.extract_EFIT_outputs(client=client, shot=shot, time_indices=time_indices)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PSI QUANTITIES\n",
    "# -------------\n",
    "psi_total_err = np.zeros(n)\n",
    "psi_total_err_med = np.zeros(n)\n",
    "psi_axis_err = np.zeros(n)\n",
    "psi_bndry_err = np.zeros(n)\n",
    "for i in range(0, n):\n",
    "    psi_total_err[i] = np.max(np.abs(freegsnke_eqs[i].psi() - efit_psi_total[i,:,:]))\n",
    "    psi_total_err_med[i] = np.median(np.abs(freegsnke_eqs[i].psi() - efit_psi_total[i,:,:]))\n",
    "    psi_axis_err[i] = np.abs(freegsnke_eqs[i].psi_axis - efit_psi_axis[i])\n",
    "    psi_bndry_err[i] = np.abs(freegsnke_eqs[i].psi_bndry - efit_psi_boundary[i])\n",
    "\n",
    "\n",
    "fig1, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 6.3), dpi=80)\n",
    "plt.subplots_adjust(hspace=0.08)  # Adjust the horizontal space between subplots\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, psi_total_err, color='navy',  linewidth=1, marker='.', markersize=4, label='FreeGSNKE', zorder=10) \n",
    "ax1.plot(times, psi_total_err_med, color='navy', linestyle='dotted', linewidth=1, marker='.', markersize=4, zorder=10) \n",
    "ax1.set_ylabel(r'$\\psi$ abs. diff')\n",
    "ax1.set_yscale('log')\n",
    "ax1.set_xticks(np.arange(0,max(times)+0.1, 0.05))\n",
    "ax1.set_yticks([1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax1.set_xticklabels([])\n",
    "ax1.set_xlim(min(times), max(times))\n",
    "ax1.set_axisbelow(True)\n",
    "\n",
    "ax2.grid(zorder=0, alpha=0.75)\n",
    "ax2.plot(times, psi_axis_err, color='navy', marker='.', markersize=4, linewidth=1, label='FreeGSNKE', zorder=10) \n",
    "ax2.set_ylabel(r'$\\psi_a$ abs. diff.')\n",
    "ax2.set_yscale('log')\n",
    "ax2.set_xticks(np.arange(0,max(times)+0.1,0.05))\n",
    "ax2.set_yticks([1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax2.set_xticklabels([])\n",
    "ax2.set_xlim(min(times), max(times))\n",
    "ax2.set_axisbelow(True)\n",
    "\n",
    "ax3.grid(zorder=0, alpha=0.75)\n",
    "ax3.plot(times, psi_bndry_err, color='navy', marker='.', markersize=4, linewidth=1, label='FreeGSNKE', zorder=10) \n",
    "ax3.set_xlabel(r'Shot time [$s$]')\n",
    "ax3.set_ylabel(r'$\\psi_b$ abs. diff.')\n",
    "ax3.set_yscale('log')\n",
    "ax3.set_xticks(np.arange(0,max(times)+0.1,0.05))\n",
    "ax3.set_yticks([1e-5, 1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax3.set_xlim(min(times), max(times))\n",
    "ax3.set_axisbelow(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Jtor \n",
    "# -------------\n",
    "jtor_total_err = np.zeros(n)\n",
    "\n",
    "for i in range(0, n):\n",
    "    jtor_total_err[i] = np.max(np.abs(freegsnke_profiles[i].jtor - efit_jtor[i,:,:]))/np.max(np.abs(efit_jtor[i,:,:]))\n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(12, 2.1), dpi=80)\n",
    "plt.subplots_adjust(hspace=0.08)  # Adjust the horizontal space between subplots\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, jtor_total_err, color='navy',  linewidth=1, marker='.', markersize=4, label='FreeGSNKE', zorder=10) \n",
    "ax1.set_xlabel(r'Shot time [$s$]')\n",
    "ax1.set_ylabel(r'$J_{tor}$ rel. diff')\n",
    "ax1.set_yscale('log')\n",
    "ax1.set_xticks(np.arange(0,max(times)+0.1, 0.05))\n",
    "ax1.set_yticks([1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax1.set_xlim(min(times), max(times))\n",
    "ax1.set_axisbelow(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SHAPE/POSITION QUANTITIES\n",
    "# -------------\n",
    "\n",
    "inner_r_err = np.zeros(n)\n",
    "outer_z_err = np.zeros(n)\n",
    "mag_r_err = np.zeros(n)\n",
    "mag_z_err = np.zeros(n)\n",
    "for i in range(0,n):\n",
    "\n",
    "    inner_r_err[i] = np.abs(freegsnke_eqs[i].innerOuterSeparatrix(0)[0] -  efit_midplane_inner_outer_radii[i,0])\n",
    "    outer_z_err[i] = np.abs(freegsnke_eqs[i].innerOuterSeparatrix(0)[1] -  efit_midplane_inner_outer_radii[i,1])\n",
    "    mag_r_err[i] = np.abs(freegsnke_eqs[i].magneticAxis()[0] -  efit_magnetic_axis[i,0])\n",
    "    mag_z_err[i] = np.abs(freegsnke_eqs[i].magneticAxis()[1] -  efit_magnetic_axis[i,1])\n",
    "\n",
    "fig2, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(6, 6), dpi=80)\n",
    "plt.subplots_adjust(wspace=0.2, hspace=0.08)\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, inner_r_err, color='navy',  linewidth=1, marker='.', markersize=2, label='FreeGSNKE', zorder=10) \n",
    "ax1.set_ylabel(r'$R_{in}$ abs. diff. [$m$]')\n",
    "ax1.set_yscale('log')\n",
    "ax1.set_xticks(np.arange(0.1,max(times)+0.1,0.2))\n",
    "ax1.set_yticks([1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax1.set_xticklabels([])\n",
    "ax1.set_xlim(min(times), max(times)+0.005)\n",
    "ax1.set_axisbelow(True)\n",
    "\n",
    "ax2.grid(zorder=0, alpha=0.75)\n",
    "ax2.plot(times, outer_z_err, color='navy',  linewidth=1, marker='.', markersize=2, label='FreeGSNKE', zorder=10) \n",
    "ax2.set_ylabel(r'$R_{out}$ abs. diff. [$m$]')\n",
    "ax2.set_yscale('log')\n",
    "ax2.set_xticks(np.arange(0.1,max(times)+0.1,0.2))\n",
    "ax2.set_yticks([1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax2.set_xticklabels([])\n",
    "ax2.set_yticklabels([])\n",
    "ax2.set_xlim(min(times), max(times)+0.005)\n",
    "ax2.set_axisbelow(True)\n",
    "\n",
    "ax3.grid(zorder=0, alpha=0.75)\n",
    "ax3.plot(times, mag_r_err, color='navy',  linewidth=1, marker='.', markersize=2, label='FreeGSNKE', zorder=10) \n",
    "# ax3.set_aspect('equal')\n",
    "ax3.set_xlabel(r'Shot time [$s$]')\n",
    "ax3.set_ylabel(r'$R_{m}$ abs. diff. [$m$]')\n",
    "ax3.set_yscale('log')\n",
    "ax3.set_xticks(np.arange(0.1,max(times)+0.1,0.2))\n",
    "ax3.set_yticks([1e-5, 1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax3.set_xlim(min(times), max(times)+0.005)\n",
    "ax3.set_axisbelow(True)\n",
    "\n",
    "ax4.grid(zorder=0, alpha=0.75)\n",
    "ax4.plot(times, mag_z_err, color='navy',  linewidth=1, marker='.', markersize=2, label='FreeGSNKE', zorder=10) \n",
    "ax4.set_xlabel(r'Shot time [$s$]')\n",
    "ax4.set_ylabel(r'$Z_{m}$ abs. diff. [$m$]')\n",
    "ax4.set_yscale('log')\n",
    "ax4.set_xticks(np.arange(0.1,max(times)+0.1,0.2))\n",
    "ax4.set_yticks([1e-5, 1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax4.set_yticklabels([])\n",
    "ax4.set_xlim(min(times), max(times)+0.005)\n",
    "ax4.set_axisbelow(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# X-POINTS \n",
    "# -------------\n",
    "# the x point calculation in EFIT doesn't seem to yield good results so we use FreeGSNKE's built-in calculator here\n",
    "\n",
    "from freegsnke import critical\n",
    "\n",
    "xpoints_freegsnke = np.zeros((n,2))\n",
    "xpoints_fiesta = np.zeros((n,2))\n",
    "xpoints_efit = np.zeros((n,2))\n",
    "\n",
    "\n",
    "box = [0.3, 0.7, -1.4, -1.0]  # box in which to search for xpoints (can change if needed)\n",
    "for i in range(0, n):\n",
    "\n",
    "    # ------------------\n",
    "    # efit x-points (calculated using FreeGSNKE function)\n",
    "    opt, xpt = critical.find_critical(eq.R, eq.Z, efit_psi_total[i,:,:])\n",
    "    if len(xpt) == 0:\n",
    "        xpoints_efit[i] = [None, None]\n",
    "    else:\n",
    "        XPT = np.array(xpt)\n",
    "        x_r = XPT[:,0]\n",
    "        x_z = XPT[:,1]\n",
    "        ind =  np.where((box[0] < x_r) & (x_r < box[1]) & (box[2] < x_z) & (x_z < box[3]))[0]\n",
    "        if len(ind) == 0:       \n",
    "            xpoints_efit[i,:] = None\n",
    "        else:\n",
    "            xpoint_r = deepcopy(x_r[ind[0]])\n",
    "            xpoint_z = deepcopy(x_z[ind[0]])\n",
    "            xpoints_efit[i,:] = np.array([xpoint_r, xpoint_z])\n",
    "\n",
    "\n",
    "    # ------------------\n",
    "    # freegsnke x-points (calculated using FreeGSNKE function)\n",
    "    x_r = freegsnke_eqs[i].xpt[:,0]\n",
    "    x_z = freegsnke_eqs[i].xpt[:,1]\n",
    "    ind = np.where((box[0] < x_r) & (x_r < box[1]) & (box[2] < x_z) & (x_z < box[3]))[0]\n",
    "    if len(ind) == 0:       \n",
    "        xpoints_freegsnke[i] = None\n",
    "    else:\n",
    "        xpoints_freegsnke[i,:] = deepcopy([x_r[ind[0]], x_z[ind[0]]])\n",
    "\n",
    "freegsnke_xpoint_err = xpoints_freegsnke - xpoints_efit\n",
    "\n",
    "fig3, (ax1) = plt.subplots(1, 1, figsize=(5, 5), dpi=80)\n",
    "plt.subplots_adjust(wspace=-0.7)  # Adjust the horizontal space between subplots\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.axvline(x=0, ymin=-1, ymax=1, linestyle='--', linewidth=1, color='k')\n",
    "ax1.axhline(y=0, xmin=-1, xmax=1, linestyle='--', linewidth=1, color='k')\n",
    "ax1.scatter(freegsnke_xpoint_err[:,0],freegsnke_xpoint_err[:,1], color='navy', s=6, marker='.', label='FreeGSNKE', zorder=10)\n",
    "# ax1.scatter(0, 0, color='r', marker='*', zorder=11) \n",
    "ax1.set_aspect('equal')\n",
    "ax1.set_xlabel(r'$R_{x}$ diff. [$m$]')\n",
    "ax1.set_ylabel(r'$Z_{x}$ diff. [$m$]')\n",
    "#ax1.set_xlim(-0.0125, 0.0125)\n",
    "#ax1.set_ylim(-0.0125, 0.0125)\n",
    "ax1.set_axisbelow(True)\n",
    "ax1.ticklabel_format(axis='both', scilimits=(0,0))\n",
    "plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# STRIKEPOINTS (lower divertor chamber)\n",
    "# -------------\n",
    "# again, we use our own calculator for finding the strikepoints\n",
    "\n",
    "strikes_freegsnke = np.zeros((n,2))\n",
    "strikes_efit = np.zeros((n,2))\n",
    "\n",
    "# which strikepoint are we looking for?\n",
    "# (lower right hand corner of domain in this case)\n",
    "r_min = 0.5\n",
    "z_max = -1.5\n",
    "for i in range(0, n):\n",
    "\n",
    "    # ------------------\n",
    "    # freegsnke strikepoints\n",
    "    \n",
    "    s = mastu_tools.find_strikepoints(R=eq.R, \n",
    "                          Z=eq.Z, \n",
    "                          psi=freegsnke_eqs[i].psi(), \n",
    "                          psi_boundary=freegsnke_eqs[i].psi_bndry, \n",
    "                          wall=np.array([eq.tokamak.wall.R,eq.tokamak.wall.Z]).T)\n",
    "    s_r = s[:,0]\n",
    "    s_z = s[:,1]\n",
    "    ind = np.where((s_r > r_min) & (s_z < z_max))[0]\n",
    "    if len(np.where((s_r > r_min) & (s_z < z_max))[0]) == 0:       \n",
    "        strikes_freegsnke[i] = None\n",
    "    else:\n",
    "        strikes_freegsnke[i,:] = np.array([s_r[ind[0]], s_z[ind[0]]])\n",
    "\n",
    "    # ------------------\n",
    "    # efit strikepoints (calculated using FreeGSNKE function)\n",
    "    efit_strikes_list = mastu_tools.find_strikepoints(R=eq.R,\n",
    "                                    Z=eq.Z,\n",
    "                                    psi=efit_psi_total[i,:,:], \n",
    "                                    psi_boundary=efit_psi_boundary[i],\n",
    "                                    wall=np.array([eq.tokamak.wall.R,eq.tokamak.wall.Z]).T)\n",
    "\n",
    "    try:\n",
    "        s_r = efit_strikes_list[:,0]\n",
    "        s_z = efit_strikes_list[:,1]\n",
    "        ind = np.where((s_r > r_min) & (s_z < z_max))[0]\n",
    "        if len(np.where((s_r > r_min) & (s_z < z_max))[0]) == 0:       \n",
    "            strikes_efit[i] = None\n",
    "        else:\n",
    "            strikes_efit[i] = np.array([s_r[ind[0]], s_z[ind[0]]])\n",
    "    except:\n",
    "        strikes_efit[i] = None\n",
    "\n",
    "\n",
    "freegsnke_strikes_err = strikes_freegsnke - strikes_efit\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig4, ((ax1), (ax2)) = plt.subplots(2, 1, figsize=(12, 4.2), dpi=80)\n",
    "plt.subplots_adjust(hspace=0.08)  # Adjust the horizontal space between subplots\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, np.abs(strikes_freegsnke[:,0] - strikes_efit[:,0]), color='navy',  linewidth=1, marker='.', markersize=4, label='FreeGSNKE', zorder=10) \n",
    "ax1.set_ylabel(r'$R_{s}$ abs. diff. [$m$]')\n",
    "ax1.set_yscale('log')\n",
    "ax1.set_xticks(np.arange(0,max(times)+0.1,0.05))\n",
    "ax1.set_yticks([1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax1.set_xlim(min(times), max(times))\n",
    "ax1.set_xticklabels([])\n",
    "ax1.set_axisbelow(True)\n",
    "\n",
    "ax2.grid(zorder=0, alpha=0.75)\n",
    "ax2.plot(times, np.abs(strikes_freegsnke[:,1] - strikes_efit[:,1]), color='navy',  linewidth=1, marker='.', markersize=4, label='FreeGSNKE', zorder=10) \n",
    "ax2.set_xlabel(r'Shot time [$s$]')\n",
    "ax2.set_ylabel(r'$Z_{s}$ abs. diff. [$m$]')\n",
    "ax2.set_yscale('log')\n",
    "ax2.set_xticks(np.arange(0,max(times)+0.1,0.05))\n",
    "ax2.set_yticks([1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1])\n",
    "ax2.set_xlim(min(times), max(times))\n",
    "ax2.set_axisbelow(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some snapshots of the lower divertor for FreeGSNKE equilibria\n",
    "snaps = [6, 46, 86, 126, 166, 206]\n",
    "\n",
    "# common contour levels\n",
    "levels = np.linspace(np.min(efit_psi_total[snaps[1],:,:]),np.max(efit_psi_total[snaps[1],:,:]),40)\n",
    "\n",
    "fig5, (ax1, ax2, ax3, ax4, ax5, ax6) = plt.subplots(1, 6, figsize=(12, 6.3), dpi=80)\n",
    "plt.subplots_adjust(wspace=0.03)  # Adjust the horizontal space between subplots\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax1.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax1.contour(eq.R, eq.Z, freegsnke_eqs[snaps[0]].psi(), levels=levels, alpha=0.25)\n",
    "ax1.contour(eq.R, eq.Z, efit_psi_total[snaps[0],:,:], levels=[efit_psi_boundary[snaps[0]]], colors='tab:gray', linestyles = 'solid')\n",
    "ax1.contour(eq.R, eq.Z, freegsnke_eqs[snaps[0]].psi(), levels=[freegsnke_eqs[snaps[0]].psi_bndry], colors='navy', linestyles = 'dotted')\n",
    "ax1.grid(alpha=0.75)\n",
    "ax1.set_aspect('equal')\n",
    "ax1.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "ax1.set_ylabel(r'Height, $Z$ $[m]$')\n",
    "# ax1.set_xlim((0/2) + 0.15, (2.15/2) + 0.15)\n",
    "ax1.set_xlim((0/1) + 0.15, 1.75)\n",
    "ax1.set_ylim(-2.25, -1.35)\n",
    "ax1.set_title(f\"$t = {times[snaps[0]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax2.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax2, show=False)\n",
    "ax2.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax2.contour(eq.R, eq.Z, freegsnke_eqs[snaps[1]].psi(), levels=levels, alpha=0.25)\n",
    "ax2.contour(eq.R, eq.Z, efit_psi_total[snaps[1],:,:], levels=[efit_psi_boundary[snaps[1]]], colors='tab:gray', linestyles = 'solid')\n",
    "ax2.contour(eq.R, eq.Z, freegsnke_eqs[snaps[1]].psi(), levels=[freegsnke_eqs[snaps[1]].psi_bndry], colors='navy', linestyles = 'dotted')\n",
    "ax2.grid(alpha=0.75)\n",
    "ax2.set_aspect('equal')\n",
    "ax2.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax2.set_ylabel(r'Height $[m]$')\n",
    "# ax1.set_xlim((0/2) + 0.15, (2.15/2) + 0.15)\n",
    "ax2.set_xlim((0/1) + 0.15, 1.75)\n",
    "ax2.set_ylim(-2.25, -1.35)\n",
    "ax2.set_yticklabels([])\n",
    "ax2.set_title(f\"$t = {times[snaps[1]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax3.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax3, show=False)\n",
    "ax3.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax3.contour(eq.R, eq.Z, freegsnke_eqs[snaps[2]].psi(), levels=levels, alpha=0.25)\n",
    "ax3.contour(eq.R, eq.Z, efit_psi_total[snaps[2],:,:], levels=[efit_psi_boundary[snaps[2]]], colors='tab:gray', linestyles = 'solid')\n",
    "ax3.contour(eq.R, eq.Z, freegsnke_eqs[snaps[2]].psi(), levels=[freegsnke_eqs[snaps[2]].psi_bndry], colors='navy', linestyles = 'dotted')\n",
    "ax3.grid(alpha=0.75)\n",
    "ax3.set_aspect('equal')\n",
    "ax3.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax3.set_ylabel(r'Height $[m]$')\n",
    "# ax1.set_xlim((0/2) + 0.15, (2.15/2) + 0.15)\n",
    "ax3.set_xlim((0/1) + 0.15, 1.75)\n",
    "ax3.set_ylim(-2.25, -1.35)\n",
    "ax3.set_yticklabels([])\n",
    "ax3.set_title(f\"$t = {times[snaps[2]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax4.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax4, show=False)\n",
    "ax4.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax4.contour(eq.R, eq.Z, freegsnke_eqs[snaps[3]].psi(), levels=levels, alpha=0.25)\n",
    "ax4.contour(eq.R, eq.Z, efit_psi_total[snaps[3],:,:], levels=[efit_psi_boundary[snaps[3]]], colors='tab:gray', linestyles = 'solid')\n",
    "ax4.contour(eq.R, eq.Z, freegsnke_eqs[snaps[3]].psi(), levels=[freegsnke_eqs[snaps[3]].psi_bndry], colors='navy', linestyles = 'dotted')\n",
    "ax4.grid(alpha=0.75)\n",
    "ax4.set_aspect('equal')\n",
    "ax4.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax4.set_ylabel(r'Height $[m]$')\n",
    "# ax1.set_xlim((0/2) + 0.15, (2.15/2) + 0.15)\n",
    "ax4.set_xlim((0/1) + 0.15, 1.75)\n",
    "ax4.set_ylim(-2.25, -1.35)\n",
    "ax4.set_yticklabels([])\n",
    "ax4.set_title(f\"$t = {times[snaps[3]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax5.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax5, show=False)\n",
    "ax5.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax5.contour(eq.R, eq.Z, freegsnke_eqs[snaps[4]].psi(), levels=levels, alpha=0.25)\n",
    "ax5.contour(eq.R, eq.Z, efit_psi_total[snaps[4],:,:], levels=[efit_psi_boundary[snaps[4]]], colors='tab:gray', linestyles = 'solid')\n",
    "ax5.contour(eq.R, eq.Z, freegsnke_eqs[snaps[4]].psi(), levels=[freegsnke_eqs[snaps[4]].psi_bndry], colors='navy', linestyles = 'dotted')\n",
    "ax5.grid(alpha=0.75)\n",
    "ax5.set_aspect('equal')\n",
    "ax5.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax5.set_ylabel(r'Height $[m]$')\n",
    "# ax5.set_xlim((0/2) + 0.15, (2.15/2) + 0.15)\n",
    "ax5.set_xlim((0/1) + 0.15, 1.75)\n",
    "ax5.set_ylim(-2.25, -1.35)\n",
    "ax5.set_yticklabels([])\n",
    "ax5.set_title(f\"$t = {times[snaps[4]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax6.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax6, show=False)\n",
    "ax6.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax6.contour(eq.R, eq.Z, freegsnke_eqs[snaps[5]].psi(), levels=levels, alpha=0.25)\n",
    "ax6.contour(eq.R, eq.Z, efit_psi_total[snaps[5],:,:], levels=[efit_psi_boundary[snaps[5]]], colors='tab:gray', linestyles = 'solid')\n",
    "ax6.contour(eq.R, eq.Z, freegsnke_eqs[snaps[5]].psi(), levels=[freegsnke_eqs[snaps[5]].psi_bndry], colors='navy', linestyles = 'dotted')\n",
    "ax6.grid(alpha=0.75)\n",
    "ax6.set_aspect('equal')\n",
    "ax6.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax6.set_ylabel(r'Height $[m]$')\n",
    "# ax6.set_xlim((0/2) + 0.15, (2.15/2) + 0.15)\n",
    "ax6.set_xlim((0/1) + 0.15, 1.75)\n",
    "ax6.set_ylim(-2.25, -1.35)\n",
    "ax6.set_yticklabels([])\n",
    "ax6.set_title(f\"$t = {times[snaps[5]]}$s\")\n",
    "\n",
    "# plt.tight_layout()  # Adjust subplots to fit into the figure area without overlapping\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SEPARATRIX COMPARISONS (core chamber only)\n",
    "# -------------\n",
    "# Here we use our own function to calculate the separatrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the metrics\n",
    "efit_separatrix = [None]*n\n",
    "freegsnke_separatrix = [None]*n\n",
    "\n",
    "separatrix_dists = np.zeros(n)\n",
    "separatrix_dists_alt = np.zeros(n)\n",
    "for i in range(0, n):\n",
    "    \n",
    "    # common magnetic axis is EFIT one\n",
    "    input_opoint = [efit_magnetic_axis[i,0], efit_magnetic_axis[i,1]]\n",
    "    \n",
    "    try:\n",
    "        efit_separatrix[i], _ = mastu_tools.Separatrix(eq.R, eq.Z, efit_psi_total[i,:,:], ntheta=360, psival=1.0, input_opoint=input_opoint)\n",
    "    except:\n",
    "        efit_separatrix[i] = None\n",
    "        \n",
    "    freegsnke_separatrix[i], _ = mastu_tools.Separatrix(eq.R, eq.Z, freegsnke_eqs[i].psi(), ntheta=360, psival=1.0, input_opoint=input_opoint)\n",
    "\n",
    "    # calculate difference (two different methods)\n",
    "    separatrix_dists[i] = mastu_tools.max_euclidean_distance(freegsnke_separatrix[i], efit_separatrix[i])\n",
    "    separatrix_dists_alt[i], _, _ = mastu_tools.separatrix_areas(freegsnke_separatrix[i], efit_separatrix[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig7, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 4.2), dpi=80)\n",
    "plt.subplots_adjust(hspace=0.08)  # Adjust the horizontal space between subplots\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.plot(times, separatrix_dists, color='navy',  linewidth=1, marker='.', markersize=4, label='FreeGSNKE', zorder=10) \n",
    "ax1.set_ylabel(r'$\\zeta$ $[m]$')\n",
    "ax1.set_yscale('log')\n",
    "ax1.set_xticks(np.arange(0.1,max(times)+0.1,0.05))\n",
    "ax1.set_yticks([1e-3, 1e-2, 1e-1])\n",
    "ax1.set_xlim(min(times), max(times))\n",
    "ax1.set_xticklabels([])\n",
    "ax1.set_axisbelow(True)\n",
    "\n",
    "# Creating the second y-axis\n",
    "ax2.grid(zorder=0, alpha=0.75)\n",
    "ax2.plot(times, separatrix_dists_alt, color='navy', linestyle=\"-\", linewidth=1, marker='.', markersize=4, zorder=10) \n",
    "ax2.set_xlabel(r'Shot time [$s$]')\n",
    "ax2.set_ylabel(r'$\\eta$')\n",
    "ax2.set_yscale('log') \n",
    "ax2.set_xticks(np.arange(0.1,max(times)+0.1,0.05))\n",
    "ax2.set_yticks([1e-3, 1e-2, 1e-1])\n",
    "ax2.set_xlim(min(times), max(times))\n",
    "ax2.set_axisbelow(True)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.path import Path\n",
    "\n",
    "# some snapshots of the lower divertor for FreeGSNKE equilibria\n",
    "snaps = [6, 46, 86, 126, 166, 206]\n",
    "\n",
    "# common contour levels\n",
    "levels = np.linspace(np.min(efit_psi_total[snaps[1],:,:]),np.max(efit_psi_total[snaps[1],:,:]),40)\n",
    "\n",
    "fig7, (ax1, ax2, ax3, ax4,ax5, ax6) = plt.subplots(1, 6, figsize=(12, 6.3), dpi=80)\n",
    "plt.subplots_adjust(wspace=0.02)  # Adjust the horizontal space between subplots\n",
    "\n",
    "# Define the clipping path using the wall coordinates\n",
    "clip_path = Path(np.column_stack((tokamak.wall.R, tokamak.wall.Z)))\n",
    "mask = clip_path.contains_points(np.column_stack((eq.R.flatten(), eq.Z.flatten())))\n",
    "mask = mask.reshape(eq.R.shape)\n",
    "# mask = np.full((65,65), True)\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax1.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax1.contour(eq.R, eq.Z, np.ma.masked_array(efit_psi_total[snaps[0],:,:], mask=~mask), levels=[efit_psi_boundary[snaps[0]]], colors='tab:gray', linestyles = 'solid', linewidths=1.4)\n",
    "ax1.contour(eq.R, eq.Z, np.ma.masked_array(freegsnke_eqs[snaps[0]].psi(), mask=~mask), levels=[freegsnke_eqs[snaps[0]].psi_bndry], colors='navy', linestyles = 'dotted', linewidths=1.0)\n",
    "ax1.grid(alpha=0.75)\n",
    "ax1.set_aspect('equal')\n",
    "ax1.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "ax1.set_ylabel(r'Height, $Z$ $[m]$')\n",
    "ax1.set_xticks(np.arange(0.5,2.1,0.5))\n",
    "ax1.set_xlim(0.15, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "ax1.set_title(f\"$t = {times[snaps[0]]}$s\")\n",
    "\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax2.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax2, show=False)\n",
    "ax2.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax2.contour(eq.R, eq.Z, np.ma.masked_array(efit_psi_total[snaps[1],:,:], mask=~mask), levels=[efit_psi_boundary[snaps[1]]], colors='tab:gray', linestyles = 'solid', linewidths=1.4)\n",
    "ax2.contour(eq.R, eq.Z, np.ma.masked_array(freegsnke_eqs[snaps[1]].psi(), mask=~mask), levels=[freegsnke_eqs[snaps[1]].psi_bndry], colors='navy', linestyles = 'dotted', linewidths=1.0)\n",
    "ax2.grid(alpha=0.75)\n",
    "ax2.set_aspect('equal')\n",
    "ax2.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "ax2.set_xticks(np.arange(0.5,2.1,0.5))\n",
    "ax2.set_xlim(0.15, 2.15)\n",
    "ax2.set_ylim(-2.25, 2.25)\n",
    "ax2.set_yticklabels([])\n",
    "ax2.set_title(f\"$t = {times[snaps[1]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax3.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax3, show=False)\n",
    "ax3.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax3.contour(eq.R, eq.Z, np.ma.masked_array(efit_psi_total[snaps[2],:,:], mask=~mask), levels=[efit_psi_boundary[snaps[2]]], colors='tab:gray', linestyles = 'solid', linewidths=1.4)\n",
    "ax3.contour(eq.R, eq.Z, np.ma.masked_array(freegsnke_eqs[snaps[2]].psi(), mask=~mask), levels=[freegsnke_eqs[snaps[2]].psi_bndry], colors='navy', linestyles = 'dotted', linewidths=1.0)\n",
    "ax3.grid(alpha=0.75)\n",
    "ax3.set_aspect('equal')\n",
    "ax3.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax3.set_ylabel(r'Height $[m]$')\n",
    "ax3.set_xticks(np.arange(0.5,2.1,0.5))\n",
    "ax3.set_xlim(0.15, 2.15)\n",
    "ax3.set_ylim(-2.25, 2.25)\n",
    "ax3.set_yticklabels([])\n",
    "ax3.set_title(f\"$t = {times[snaps[2]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax4.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax4, show=False)\n",
    "ax4.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax4.contour(eq.R, eq.Z, np.ma.masked_array(efit_psi_total[snaps[3],:,:], mask=~mask), levels=[efit_psi_boundary[snaps[3]]], colors='tab:gray', linestyles = 'solid', linewidths=1.4)\n",
    "ax4.contour(eq.R, eq.Z, np.ma.masked_array(freegsnke_eqs[snaps[3]].psi(), mask=~mask), levels=[freegsnke_eqs[snaps[3]].psi_bndry], colors='navy', linestyles = 'dotted', linewidths=1.0)\n",
    "ax4.grid(alpha=0.75)\n",
    "ax4.set_aspect('equal')\n",
    "ax4.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax4.set_ylabel(r'Height $[m]$')\n",
    "ax4.set_xticks(np.arange(0.5,2.1,0.5))\n",
    "ax4.set_xlim(0.15, 2.15)\n",
    "ax4.set_ylim(-2.25, 2.25)\n",
    "ax4.set_yticklabels([])\n",
    "ax4.set_title(f\"$t = {times[snaps[3]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax5.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax5, show=False)\n",
    "ax5.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax5.contour(eq.R, eq.Z, np.ma.masked_array(efit_psi_total[snaps[4],:,:], mask=~mask), levels=[efit_psi_boundary[snaps[4]]], colors='tab:gray', linestyles = 'solid', linewidths=1.4)\n",
    "ax5.contour(eq.R, eq.Z, np.ma.masked_array(freegsnke_eqs[snaps[4]].psi(), mask=~mask), levels=[freegsnke_eqs[snaps[4]].psi_bndry], colors='navy', linestyles = 'dotted', linewidths=1.0)\n",
    "ax5.grid(alpha=0.75)\n",
    "ax5.set_aspect('equal')\n",
    "ax5.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax5.set_ylabel(r'Height $[m]$')\n",
    "ax5.set_xticks(np.arange(0.5,2.1,0.5))\n",
    "ax5.set_xlim(0.15, 2.15)\n",
    "ax5.set_ylim(-2.25, 2.25)\n",
    "ax5.set_yticklabels([])\n",
    "ax5.set_title(f\"$t = {times[snaps[4]]}$s\")\n",
    "\n",
    "rect0 = patches.Rectangle((0.26, -2.21), 1.74, 4.42, edgecolor='none', facecolor='lightgrey', alpha=0.4, zorder=0)\n",
    "ax6.add_patch(rect0)\n",
    "eq.tokamak.plot(axis=ax6, show=False)\n",
    "ax6.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)\n",
    "ax6.contour(eq.R, eq.Z, np.ma.masked_array(efit_psi_total[snaps[5],:,:], mask=~mask), levels=[efit_psi_boundary[snaps[5]]], colors='tab:gray', linestyles = 'solid', linewidths=1.4)\n",
    "ax6.contour(eq.R, eq.Z, np.ma.masked_array(freegsnke_eqs[snaps[5]].psi(), mask=~mask), levels=[freegsnke_eqs[snaps[5]].psi_bndry], colors='navy', linestyles = 'dotted', linewidths=1.0)\n",
    "ax6.grid(alpha=0.75)\n",
    "ax6.set_aspect('equal')\n",
    "ax6.set_xlabel(r'Major radius, $R$ $[m]$')\n",
    "# ax6.set_ylabel(r'Height $[m]$')\n",
    "ax6.set_xticks(np.arange(0.5,2.1,0.5))\n",
    "ax6.set_xlim(0.15, 2.15)\n",
    "ax6.set_ylim(-2.25, 2.25)\n",
    "ax6.set_yticklabels([])\n",
    "ax6.set_title(f\"$t = {times[snaps[5]]}$s\")\n",
    "\n",
    "# plt.tight_layout()  # Adjust subplots to fit into the figure area without overlapping\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# INITIALISE PROBES\n",
    "# -------------\n",
    "eq.tokamak.probes.initialise_setup(eq) \n",
    "print('Magnetic probes initialised in the tokamak.')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CALCULATE PROBE VALUES IN FREEGSNKE\n",
    "# ------------- \n",
    "\n",
    "# get the names of the probes\n",
    "probe_names = eq.tokamak.probes.floop_order + eq.tokamak.probes.pickup_order\n",
    "        \n",
    "probe_values = np.zeros((len(probe_names),len(times)))\n",
    "for i in range(0,n):\n",
    "    \n",
    "    # compute fluxloop and pickup coil values from equilibrium \n",
    "    probe_values[:,i] = np.concatenate([\n",
    "        tokamak.probes.calculate_fluxloop_value(freegsnke_eqs[i])*2*np.pi, # multiply by 2pi to match flux units in UDA\n",
    "        tokamak.probes.calculate_pickup_value(freegsnke_eqs[i])\n",
    "    ])\n",
    "    \n",
    "# fluxloop and pickup coil locations\n",
    "probe_locations = np.concatenate([\n",
    "    eq.tokamak.probes.floop_pos, \n",
    "    eq.tokamak.probes.pickup_pos[:,0::2]\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EXTRACT FLUXLOOP/PICKUP COIL MEASUREMENTS FROM EFIT DATA AT EACH TIME STEP REQUIRED\n",
    "# -------------\n",
    "\n",
    "# extract the probe names with values (will be less than the names in probe_names)\n",
    "mastu_names = list(efit_fluxloops['names']) + list(efit_pickups['names'])\n",
    "\n",
    "# concatenate all values (not all probes will return values)\n",
    "# mastu_values = np.concatenate((raw_fluxloops['target'].T,raw_pickups['target'].T), axis=0)\n",
    "mastu_values = np.concatenate((efit_fluxloops['computed'].T, efit_pickups['computed'].T), axis=0)\n",
    "mastu_values = mastu_values[:,time_indices]\n",
    "\n",
    "mastu_sigmas = np.concatenate((efit_fluxloops['sigmas'].T, efit_pickups['sigmas'].T), axis=0)\n",
    "mastu_sigmas = mastu_sigmas[:,time_indices]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# find the correct indices to link the measurement and FreeGSNKE data\n",
    "ordered_freegsnke_values = np.zeros((len(mastu_names),len(times)))\n",
    "ordered_freegsnke_locations = np.zeros((len(mastu_names),2))\n",
    "for i, name in enumerate(mastu_names):\n",
    "    index = probe_names.index(name)  \n",
    "    ordered_freegsnke_values[i,:] = probe_values[index,:]\n",
    "    ordered_freegsnke_locations[i,:] = probe_locations[index]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PLOT THE ABS DIFFERENCES\n",
    "# ------------- \n",
    "\n",
    "diff = np.abs(mastu_values - ordered_freegsnke_values)\n",
    "i = len(efit_fluxloops['names'])\n",
    "\n",
    "\n",
    "fig7, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 4.2), dpi=125)\n",
    "plt.subplots_adjust(hspace=0.08)  # Adjust the horizontal space between subplots\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.5)\n",
    "for j in range(0,i):\n",
    "    ax1.plot(times, diff[j,:], linestyle=\"-\", linewidth=0.75, alpha=0.7)\n",
    "ax1.plot(times, np.max(mastu_sigmas[0:i,:], axis=0), 'k', linestyle=\"--\", linewidth=1.2)\n",
    "ax1.text(0.08, 5e-8, rf'{i}/{len(eq.tokamak.probes.floop_order)} fluxloops', color='k', fontsize=8, bbox=dict(facecolor='white', edgecolor='black', boxstyle='round,pad=0.5'))\n",
    "\n",
    "# ax1.set_aspect('equal')\n",
    "ax1.set_xlabel(r'Shot time [$s$]')\n",
    "ax1.set_ylabel(r'$\\psi$ abs. diff.')\n",
    "ax1.set_yscale('log')\n",
    "# ax1.set_ylim(-2.2, 2.2)\n",
    "ax1.set_xticks(np.arange(0.1,max(times)+0.1,0.05))\n",
    "ax1.set_yticks([1e-8, 1e-6, 1e-4, 1e-2, 1e0])\n",
    "ax1.set_xticklabels([])\n",
    "ax1.set_xlim(min(times), max(times))\n",
    "ax1.set_ylim(1e-8, 1e0)\n",
    "# ax1.legend(loc='lower left')\n",
    "ax1.set_axisbelow(True)\n",
    "\n",
    "# Creating the second y-axis\n",
    "ax2.grid(zorder=0, alpha=0.5)\n",
    "for j in range(i,len(mastu_names)):\n",
    "    ax2.plot(times, diff[j,:], linestyle=\"-\", linewidth=0.75, alpha=0.7)\n",
    "ax2.plot(times, np.max(mastu_sigmas[i::,:], axis=0), 'k', linestyle=\"--\", linewidth=1.2)\n",
    "ax2.text(0.08, 5e-8, rf'{(len(mastu_names)-i)}/{len(eq.tokamak.probes.pickup_order)} pickup coils', color='k', fontsize=8, bbox=dict(facecolor='white', edgecolor='black', boxstyle='round,pad=0.5'))\n",
    "\n",
    "ax2.set_xlabel(r'Shot time [$s$]')\n",
    "ax2.set_ylabel(r'$\\mathbf{B} \\cdot \\mathbf{\\hat{n}}$ abs. diff.')\n",
    "ax2.set_yscale('log') \n",
    "ax2.set_xticks(np.arange(0.1,max(times)+0.1,0.05))\n",
    "ax2.set_yticks([1e-8, 1e-6, 1e-4, 1e-2, 1e0])\n",
    "ax2.set_ylim(1e-8, 1e0)\n",
    "ax2.set_xlim(min(times), max(times))\n",
    "ax2.set_axisbelow(True)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "freegsnke_equilibrium",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}